import matplotlib.pyplot as plt
import tables as tb
import numpy as np
from cvxopt import matrix,solvers,lapack
import os
from scipy.linalg import toeplitz
import help_func

def curvature(alpha, smooth, chisqr):
        smooth = np.log(np.array(smooth))
        chisqr = np.log(np.array(chisqr))
        alpha = np.array(alpha)
        # 1st D
        dalpha = np.diff(alpha)
        dsmooth = np.diff(smooth) / dalpha
        dchisqr = np.diff(chisqr) / dalpha
        dalpha = dalpha[0:-1]
        # 2ed D
        d2smooth = np.diff(dsmooth) / dalpha
        d2chisqr = np.diff(dchisqr) / dalpha
        dsmooth = dsmooth[0:-1]
        dchisqr = dchisqr[0:-1]
        # curvature
        k = np.abs((dsmooth * d2chisqr - dchisqr * d2smooth) / np.power(dchisqr * dchisqr + d2smooth * d2smooth, 3. / 2.))
        return {"alpha":alpha[0:-2], "k":k, "maxk":k.argmax()}
class ACLap:
    def __init__(self, simData,mxdata):
        self.alphas=np.logspace(4,6,mxdata.numAlpha)
        print self.alphas
        diff=np.zeros_like(mxdata.omegaP)
        diff[0]=2.
        diff[-1]=-1.
        diff[1]=-1.
        D=toeplitz(diff)
        ProjDiag=np.ones_like(mxdata.omegaP)
        ProjDiag[0:2]=0;
        Proj=np.diag(ProjDiag)
        D=np.dot(np.dot(Proj,D),Proj)
        dw=mxdata.omegaP[1]-mxdata.omegaP[0]
        D=D/(dw**2)
        self.D=D
        solvers.options['feastol']=1e-8
        solvers.options['abstol']=1e-8
        solvers.options['reltol']=1e-8
        self.q=-matrix(np.dot(simData.kp.T,simData.datap).T)
        dG=np.ones_like(mxdata.omegaP)
        self.G=matrix(-np.diag(dG))
        self.h=matrix(np.zeros_like(mxdata.omegaP))
        self.ktk=np.dot(simData.kp.T,simData.kp)
        self.simData=simData
        self.beta=mxdata.beta
        self.mxData=mxdata
    def solve(self,alpha):
        res=help_func.LapResult()
        P=matrix(self.ktk+alpha*self.D)#np.dot(D.T,D))
        sol=solvers.coneqp(P,self.q, self.G, self.h)
        diff=(np.dot(self.simData.kp,sol['x']).T-self.simData.datap)
        res.chisqr=((np.dot(diff,diff.T))/(self.beta/2.))[0][0]
        res.smooth=(np.dot(sol['x'].T,np.dot(self.D,sol['x'])))[0][0]
        res.f=sol['x']
        return res
    def run(self):
        allRes=[self.solve(alpha) for alpha in self.alphas]
        return allRes
    def printAll(self,allRes):
	dw=self.mxData.omegaP[1]-self.mxData.omegaP[0]
        fig1 = plt.figure()
        ax1 = fig1.add_subplot(111)
        for res in allRes:
            ax1.plot(self.mxData.omegaP, res.f/dw)
	ax1.set_ylim(0.,.1)
        fig1.savefig(self.mxData.dump_dir + "/LapAllA.pdf")
        fig1.clf()
        smooth=np.array([res.smooth for res in allRes])
        chisqr=np.array([res.chisqr for res in allRes])
	print chisqr
        cur=curvature(self.alphas,smooth,chisqr)
        ax1 = fig1.add_subplot(111)
        ax1.plot(cur['alpha'],cur['k'])
        fig1.savefig(self.mxData.dump_dir +"/curvature.pdf")
        fig1.clf()
        ax1 = fig1.add_subplot(111)
        ax1.plot(self.mxData.omegaP,allRes[cur['maxk']].f/dw)
        maxo=8
        ax1.set_xlim(0, maxo)
        ax1.set_ylim(0.,.1)
        fig1.savefig(self.mxData.dump_dir +"/maxcurvatureLap1.pdf")
        maxo=1
        ax1.set_xlim(0, maxo)
        ax1.set_ylim(0.,.1)
        fig1.savefig(self.mxData.dump_dir +"/maxcurvatureLap2.pdf")

